{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "import argparse\n",
    "import os\n",
    "import torch\n",
    "from collections import OrderedDict\n",
    "import math\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import scipy.io\n",
    "from scipy.interpolate import griddata\n",
    "from mpl_toolkits.axes_grid1 import make_axes_locatable\n",
    "import matplotlib.gridspec as gridspec\n",
    "from torch.utils.data import DataLoader\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch\n",
    "import seaborn as sns\n",
    "import pylab as py\n",
    "import time\n",
    "from pyDOE import lhs\n",
    "import warnings\n",
    "\n",
    "sys.path.insert(0, '../../../Scripts/')\n",
    "from models_pde import Generator, Discriminator, Q_Net\n",
    "\n",
    "from pid import *\n",
    "\n",
    "warnings.filterwarnings('ignore')\n",
    "\n",
    "np.random.seed(1234)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Hyper-parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_epochs = 30000\n",
    "lambda_prob = 0.1\n",
    "lambda_q = 0.5\n",
    "noise = 0.1\n",
    "\n",
    "g_hid_dim = 50\n",
    "g_num_layer = 4\n",
    "\n",
    "k_hid_dim = 50\n",
    "k_num_layer = 4\n",
    "\n",
    "d_hid_dim = 50 \n",
    "d_num_layer = 3\n",
    "\n",
    "q_hid_dim = 50\n",
    "q_num_layer = 4"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### CUDA"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# CUDA support \n",
    "if torch.cuda.is_available():\n",
    "    device = torch.device('cuda:4')\n",
    "else:\n",
    "    device = torch.device('cpu')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = np.load('../../../datasets/nonlinear2d_data.npz')\n",
    "X = data['X']\n",
    "K = data['k']\n",
    "U = data['u']\n",
    "\n",
    "N = 10000\n",
    "N_f = N\n",
    "N_u = 200\n",
    "N_b = 100 # for one boundary\n",
    "q = 1.\n",
    "u_0 = - 10.\n",
    "ksat = 10.\n",
    "\n",
    "X_dim = 1 \n",
    "Y_dim = 1\n",
    "Z_dim = 2\n",
    "\n",
    "L1 = 10.\n",
    "L2 = 10.\n",
    "\n",
    "\n",
    "X_u = np.zeros((N_u,2))\n",
    "Y_u = np.zeros((N_u,1))\n",
    "X_f = np.zeros((N_f,2))\n",
    "\n",
    "# Boundary points\n",
    "x1_b1 = np.zeros(N_b)[:,None]\n",
    "x2_b1 = L2 * np.random.random(N_b)[:,None]\n",
    "X_b1 = np.hstack((x1_b1, x2_b1))\n",
    "\n",
    "x1_b2 = L1 * np.random.random(N_b)[:,None]\n",
    "x2_b2 = np.zeros(N_b)[:,None]\n",
    "X_b2 = np.hstack((x1_b2, x2_b2))\n",
    "\n",
    "\n",
    "x1_b3 = L1 * np.ones(N_b)[:,None]\n",
    "x2_b3 = L2 * np.random.random(N_b)[:,None]\n",
    "X_b3 = np.hstack((x1_b3, x2_b3))\n",
    "\n",
    "\n",
    "x1_b4 = L1 * np.random.random(N_b)[:,None]\n",
    "x2_b4 = L2 * np.ones(N_b)[:,None]\n",
    "X_b4 = np.hstack((x1_b4, x2_b4))\n",
    "\n",
    "\n",
    "X_b = np.hstack((X_b1, X_b2))\n",
    "X_b = np.hstack((X_b, X_b3))\n",
    "X_b = np.hstack((X_b, X_b4))\n",
    "\n",
    "\n",
    "# Collocation points\n",
    "X1_f = L1 * np.random.random(N_f)[:,None]\n",
    "X2_f = L2 * np.random.random(N_f)[:,None]\n",
    "X_f = np.hstack((X1_f, X2_f))\n",
    "\n",
    "U_data = U\n",
    "X_data = X\n",
    "\n",
    "idx_u = np.random.choice(N, N_u, replace=False)\n",
    "for i in range(N_u):\n",
    "    X_u[i,:] = X_data[idx_u[i],:]\n",
    "    Y_u[i,:] = U_data[idx_u[i]]\n",
    "    \n",
    "# Normalize data\n",
    "lb = np.array([0.0, 0.0])\n",
    "ub = np.array([10.0, 10.0])\n",
    "lbb = np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0])\n",
    "ubb = np.array([10.0, 10.0, 10.0, 10.0, 10.0, 10.0, 10.0, 10.0])\n",
    "X_u = (X_u - lb) - 0.5*(ub - lb)\n",
    "X_b = (X_b - lbb) - 0.5*(ubb - lbb)\n",
    "X_f = (X_f - lb) - 0.5*(ub - lb)\n",
    "\n",
    "\n",
    "x1_u = X_u[:,0:1]   # dimension  N_u x 1\n",
    "x2_u = X_u[:,1:2]   # dimension  N_u x 1\n",
    "y_u = Y_u           # dimension N_u\n",
    "\n",
    "x1_f = X_f[:,0:1]   # dimension N_f x 1\n",
    "x2_f = X_f[:,1:2]   # dimension N_f x 1\n",
    "\n",
    "# Position of the boundary \n",
    "x1_b1 = X_b[:,0:1]\n",
    "x2_b1 = X_b[:,1:2]\n",
    "x1_b2 = X_b[:,2:3]\n",
    "x2_b2 = X_b[:,3:4]\n",
    "x1_b3 = X_b[:,4:5]\n",
    "x2_b3 = X_b[:,5:6]\n",
    "x1_b4 = X_b[:,6:7]\n",
    "x2_b4 = X_b[:,7:8]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load Models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "D = Discriminator(in_dim = 4, out_dim = 1, \n",
    "                  hid_dim = d_hid_dim, \n",
    "                  num_layers = d_num_layer\n",
    "                 ).to(device)\n",
    "\n",
    "G = Generator(in_dim = 4, out_dim = 1, \n",
    "              hid_dim = g_hid_dim, \n",
    "              num_layers = g_num_layer\n",
    "             ).to(device)\n",
    "\n",
    "k_G = Generator(in_dim = 1, out_dim = 1, \n",
    "              hid_dim = k_hid_dim, \n",
    "              num_layers = k_num_layer\n",
    "             ).to(device)\n",
    "\n",
    "Q = Q_Net(in_dim = 3, out_dim = 2, \n",
    "          hid_dim = q_hid_dim, \n",
    "          num_layers = q_num_layer\n",
    "         ).to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lambdas = [lambda_prob, lambda_q]\n",
    "pid = Darcy_PID(X_u, Y_u, X_b, X_f, N_u, \n",
    "                  G, k_G, D, Q, device, num_epochs, \n",
    "                  lambdas, noise\n",
    "                 )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pid.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X_star = X\n",
    "u_star = U.T\n",
    "k_star = K.T\n",
    "ksat = 10.\n",
    "\n",
    "X_star_norm = (X_star - lb) - 0.5*(ub - lb)\n",
    "u_pred_list = []\n",
    "f_pred_list = []\n",
    "k_pred_list = []\n",
    "for run in range(500):\n",
    "    u_pred, f_pred, k_pred = pid.predict(X_star_norm)\n",
    "    k_pred /= ksat\n",
    "    u_pred_list.append(u_pred)\n",
    "    f_pred_list.append(f_pred)\n",
    "    k_pred_list.append(k_pred)\n",
    "\n",
    "    \n",
    "u_pred_arr = np.array(u_pred_list)\n",
    "f_pred_arr = np.array(f_pred_list)\n",
    "k_pred_arr = np.array(k_pred_list)\n",
    "\n",
    "uuu_mu_pred = u_pred_arr.mean(axis=0)\n",
    "fff_mu_pred = f_pred_arr.mean(axis=0)\n",
    "uuu_Sigma_pred = u_pred_arr.var(axis=0)\n",
    "fff_Sigma_pred = f_pred_arr.var(axis=0)\n",
    "kkk_mu_pred = k_pred_arr.mean(axis=0)\n",
    "kkk_Sigma_pred = k_pred_arr.var(axis=0)\n",
    "\n",
    "error_u = np.linalg.norm(u_star.T-uuu_mu_pred,2)/np.linalg.norm(u_star.T,2)\n",
    "error_k = np.linalg.norm(k_star.T-kkk_mu_pred,2)/np.linalg.norm(k_star.T,2)\n",
    "print('Error u: %e' % (error_u))                     \n",
    "print('Residual: %e' % (f_pred**2).mean())\n",
    "print('Error k: %e' % (error_k))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
